Repository navigation
Conversation
…han inputs utils._convert_to_result paired inputs and predictions with zip(), which stops at the shorter of the two, so inputs without a prediction were dropped without an error. It now raises a ValueError that gives the counts and names any dict output without one entry per input. Extra predictions are still ignored, so pipelines in which every input already gets a result are unaffected. Also fixes three tests that only passed because of the truncation. Fixes apache#40928
There was a problem hiding this comment.
📋 Review Summary
This PR updates utils._convert_to_result to raise a descriptive ValueError when a model handler produces fewer predictions than inputs in a batch, preventing silent data loss caused by zip() truncation. It also updates three existing unit tests that inadvertently relied on this truncation behavior and adds comprehensive test coverage for the new validation.
🔍 General Feedback
- The implementation in
utils.pycleanly handles both sequence/generator and dictionary prediction outputs, providing clear, actionable error messages (including identifying specific short keys in dictionary outputs). - Keeping extra predictions ignored preserves backward compatibility for existing single-element batch pipelines while catching silent data drops.
- The test additions in
utils_test.pyand fixes insklearn_inference_test.pyandtensorflow_inference_test.pyare thorough and well-structured.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40929 +/- ##
==========================================
+ Coverage 79.83% 80.09% +0.26%
==========================================
Files 494 494
Lines 83575 83628 +53
==========================================
+ Hits 66722 66985 +263
+ Misses 16853 16643 -210
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Assigning reviewers: R: @claudevdm for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
Fixes #40928
utils._convert_to_resultpairs each batch of inputs with the model's predictions usingzip(), which stops at the shorter of the two. When a handler produced fewer predictions than inputs, the leftover inputs got no result and no error, andnum_inferencesstill counted them. That's how #36587 and #40913 went unnoticed, and it still happens with built-in handlers, for example with PyTorch or Keras models that return a tuple of outputs, and with Hugging Face causal LMs that returnpast_key_values(details in the issue).This PR:
_convert_to_resultraise aValueErrorwhen there are fewer predictions than inputs. The message gives the counts, names any dict output that doesn't have one entry per input, and suggests a fix (an inference function that returns one prediction per input, orinference_args={'use_cache': False}for Hugging Face models);sklearn_inference_test.py,test_custom_inference_fnused an inference function that returned one prediction for three inputs, and itszip()-based check compared only the first result. Intensorflow_inference_test.py, the two env-var tests mocked the model, so their pipelines produced no results at all. The fakes now return one prediction per input, and the scikit-learn test also checks the number of results;with_exception_handling(), and a CHANGES.md entry.Testing:
utils_test.py. The 4 that check the new error fail without this change; the other 5 cover behavior that doesn't change.apache_beam/mltest suite with torch, TensorFlow, transformers, ONNX, scikit-learn and XGBoost installed (Python 3.12): 609 passed, against 600 on master. The only differences are the 9 new tests. The tests that can't run locally (they need GCP credentials or Docker) fail the same way on both.The issue also asks whether RunInference should check the count itself, to cover custom handlers that pair results with their own
zip(). I left that out to keep this change to the shared helper. If you'd prefer a warning instead of an error, or the check in RunInference itself, I'm happy to change it.R: @jrmccluskey